//! Tracks scheduled messages in near-future ticks.
//!
//! Really naive implementation.

use std::collections::HashMap;

use craftdig_datamodel::{
    chunk::{BlockMap, ChunkMap},
    message::{MessageContainer, MsgTyId},
};
use craftdig_math::*;

/// Tracker for the pending ticks in a chunk.
pub type ChunkSched = ChunkMap<Vec<MessageContainer>>;

#[derive(Clone)]
pub struct MsgScheduler {
    ticks: HashMap<u64, TickSched>,
}

impl MsgScheduler {
    pub fn new_empty() -> Self {
        Self {
            ticks: HashMap::new(),
        }
    }

    fn prep_tick(&mut self, tick: u64) {
        self.ticks.entry(tick).or_insert_with(TickSched::new_empty);
    }

    /// Schedules a message to be delivered at some tick.  If this is for a
    /// previous tick, then we might leak data.
    pub fn schedule_msg(&mut self, tick: u64, pos: BlkCoord, mc: MessageContainer) {
        self.prep_tick(tick);
        let tsched = self.ticks.get_mut(&tick).expect("msgsched: missing tick");
        tsched.insert_msg(pos, mc);
    }

    /// Remove the full tick message entries.
    pub fn remove_tick(&mut self, tick: u64) -> Option<TickSched> {
        self.ticks.remove(&tick)
    }

    /// Removes a tick from the scheduler, returning an empty blockmap if there wasn't any.
    // TODO maybe convert to `Option`?
    pub fn take_tick_msgs(&mut self, tick: u64, mty: MsgTyId) -> BlockMap<Vec<MessageContainer>> {
        let Some(tick) = self.ticks.get_mut(&tick) else {
            return BlockMap::new();
        };

        tick.take_msgs_by_ty(mty)
    }

    /// Clears all scheduled messages for the block.
    pub fn clear_block(&mut self, block: BlkCoord) {
        self.ticks
            .iter_mut()
            .for_each(|(_, ts)| ts.clear_block(block));
    }

    /// Clears all scheduled messages.
    pub fn clear_all(&mut self) {
        self.ticks.clear();
    }

    /// Clears any tick slots before a given index.  This shouldn't be necessary with correct usage.
    pub fn clear_old_ticks_before(&mut self, before: u64) {
        self.ticks.retain(|v, _| *v >= before);
    }
}

/// Tracks messages scheduled to be delivered on a tick.
#[derive(Clone)]
pub struct TickSched {
    msg_maps: HashMap<MsgTyId, BlockMap<Vec<MessageContainer>>>,
}

impl TickSched {
    pub fn new_empty() -> Self {
        Self {
            msg_maps: HashMap::new(),
        }
    }

    /// Gets or creates a map for a certain message type, returning a mut ref.
    fn get_or_create_map(&mut self, mty: MsgTyId) -> &mut BlockMap<Vec<MessageContainer>> {
        self.msg_maps.entry(mty).or_insert_with(BlockMap::new)
    }

    /// Inserts a message for a certain block.
    pub fn insert_msg(&mut self, pos: BlkCoord, mc: MessageContainer) {
        let (cc, scp) = pos.split();
        let chunks = self.get_or_create_map(mc.ty());
        let chunk = chunks.get_or_create_chunk_mut(cc);
        let blk = chunk.get_or_create_with(scp, Vec::new);
        blk.push(mc);
    }

    /// Takes all of the messages of a certain type.
    pub fn take_msgs_by_ty(&mut self, mty: MsgTyId) -> BlockMap<Vec<MessageContainer>> {
        self.msg_maps.remove(&mty).unwrap_or_default()
    }

    /// Clears any messages for the block.
    pub fn clear_block(&mut self, pos: BlkCoord) {
        self.msg_maps.values_mut().for_each(|map| {
            map.remove(pos);
        });
    }

    /// Iterates over *all* of the messages on the tick.
    pub fn iter_messages(&self) -> impl Iterator<Item = (BlkCoord, &MessageContainer)> {
        self.msg_maps.values().flat_map(|map| {
            map.iter_entries()
                .flat_map(|(bc, mvec)| mvec.iter().map(move |mc| (bc, mc)))
        })
    }
}
